import sys
import cv2
from rtmlib import Body
import numpy as np
import os
import json
from tqdm import *
from motionbert.mb_inferencer import motionbert_inferencer

from inference_2d import inference_2d_and_save
from inference_3d import inference_3d_and_save

from utils import *

def inference_2d_3d_and_save(video_path, inferencer_2d, inferencer_3d, output_vis=True, silent=False):
    assert os.path.exists(video_path)

    reault_2d_data_path = emohugo_path_shift(video_path, "video", "pose_2d")
    reault_3d_data_path = emohugo_path_shift(reault_2d_data_path, "pose_2d", "pose_3d")
    if output_vis:
        result_2d_vis_path = emohugo_path_shift(video_path, "video", "pose_2d_visual")
        result_3d_vis_path = emohugo_path_shift(reault_2d_data_path, "pose_2d", "pose_3d_visual")
    else:
        result_2d_vis_path = result_3d_vis_path = None
    inference_2d_and_save(video_path, inferencer_2d, reault_2d_data_path, result_2d_vis_path, silent)
    inference_3d_and_save(reault_2d_data_path, inferencer_3d, reault_3d_data_path, result_3d_vis_path, block_lower=True, silent=silent)    

if __name__ == "__main__":
    inferencer_2d = Body(
        pose='rtmo', # 使用 rtmpose 和 yolo detector 会报错
        to_openpose=False,
        mode='performance',  # balanced, performance, lightweight
        backend='onnxruntime',
        device='cpu')  # onnxruntime 最高支持cuda11.8
    
    config_path = os.path.join(get_project_root(), "motionbert/configs/pose3d/MB_train_h36m.yaml")
    ckpt_path = os.path.join(get_project_root(),
                             "motionbert/checkpoint/pose3d/MB_train_h36m_dual_loss/finish_at_2024_4_14/best_epoch.bin")
    inferencer_3d = motionbert_inferencer(config_path, ckpt_path)

    video_path = os.path.join(get_project_root() , "datasets/emohugo_video/surprise/000003-S003-20.avi")

    inference_2d_3d_and_save(video_path, inferencer_2d, inferencer_3d, output_vis=True, silent=True)

    

